# 此文件为工具函数或通用函数
import shutil  # 文件操作模块
import os  # 系统操作模块
from PIL import Image  # 图像处理模块
import numpy as np  # 数组和数值运算模块

# 保存单张图片
def save_image(image_tensor, out_name):
    """
    保存单张图片
    save a single image
    :param image_tensor: torch tensor with size=(3, h, w)
    :param out_name: path+name+".jpg"
    :return: None
    """
    # print("in size", image_tensor.size())
    if len(image_tensor.size()) == 3: # 检查张量尺寸是否为3
        image_numpy = image_tensor.cpu().detach().numpy().squeeze(0)  # 转换为NumPy数组，并处理维度
        image_numpy[image_numpy > 0.5] = 255  # 设定大于0.5的像素为255
        image_numpy[image_numpy < 0.5] = 0    # 设定小于0.5的像素为0
        image_numpy = image_numpy.astype(np.uint8) # 转换数据类型为无符号整型8位
        image = Image.fromarray(image_numpy)  # 从NumPy数组创建图像
        image.save(out_name)  # 保存图像到指定路径
    else:
        raise ValueError("input tensor not with size (3, h, w)")  # 若尺寸不对，抛出异常
    return None

# 检查并创建目录
def check_mk_dir(dir_name):
    if not os.path.exists(dir_name):  # 检查目录是否存在
        os.makedirs(dir_name)  # 若不存在，创建目录

# 拆分文件
def split(ori_dir, tar_dir):
    files = os.listdir(ori_dir)  # 获取原始目录下的文件列表
    for file in files:
        if len(file) > 10:  # 检查文件名长度是否大于10
            shutil.copy(os.path.join(ori_dir, file), os.path.join(tar_dir, file))  # 复制文件到目标目录

# 创建子目录
def make_sub_dirs(base_dir, sub_dirs):
    """
    six
    :param base_dir:
    :param sub_dirs:
    :return:
    """
    check_mk_dir(base_dir)  # 创建基础目录
    for sub_dir in sub_dirs:  # 遍历子目录列表
        check_mk_dir(os.path.join(base_dir, sub_dir))  # 在基础目录下创建子目录

# 创建项目目录结构
def make_project_dir(train_dir, val_dir):
    make_sub_dirs(train_dir, ["train_images", "pth", "loss"])  # 在训练目录下创建子目录
    make_sub_dirs(val_dir, ["val_images"])  # 在验证目录下创建子目录

if __name__ == "__main__":
    pass
